package wsz.redisflowcontrol.controller;

import lombok.extern.slf4j.Slf4j;
import org.redisson.Redisson;
import org.redisson.api.RLock;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.util.ObjectUtils;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.Arrays;
import java.util.Random;
import java.util.UUID;
import java.util.concurrent.TimeUnit;

/**
 * 限流测试模拟
 * @author wsz
 * @date 2022/3/25 12:15
 **/
@Slf4j
@RestController
public class RedisFlowControl {
    
    @Autowired
    private StringRedisTemplate redisTemplate;

    @Autowired
    private Redisson redisson;

    private Random random = new Random();
    private static int total = 50;
    private static int seconds = 60;

    private int getSleep() {
        return random.nextInt(5) * 1000;
    }

    /**
     * 1.普通的分步骤判断
     */
    private String common_key = "common_key";
    @RequestMapping("common")
    public String common() {
        boolean requestFlag = false;
        try {
            // 1.前置判断
            int oldValueInt = getRequestValue(common_key);
            if (++oldValueInt > total) {
                return "0";
            }

            // 2.更新
            redisTemplate.opsForValue().set(common_key, String.valueOf(oldValueInt));
            redisTemplate.expire(common_key, seconds, TimeUnit.SECONDS);
            requestFlag = true;

            // doSomething
            int sleep = getSleep();
            Thread.sleep(sleep);
            return "1";
        } catch (Exception ex) {
            ex.printStackTrace();
        } finally {
            // 3.释放
            if (requestFlag) {
                redisTemplate.opsForValue().decrement(common_key);
            }
        }

        return "0";
    }

    /**
     * 2.lua模式
     */
    private String lua_key = "lua_key";
    @RequestMapping("lua")
    public String lua() {

        Long execute = null;
        try {
            String script =
                    "if redis.call('EXISTS', KEYS[1]) == 1 then " +
                            "if tonumber(redis.call('get', KEYS[1])) < tonumber(ARGV[1]) then " +
                                "redis.call('incr', KEYS[1]) redis.call('expire', KEYS[1], ARGV[2]) return 1 " +
                            "else return 0 " +
                            "end " +
                    "else " +
                            "redis.call('incr', KEYS[1]) redis.call('expire', KEYS[1], ARGV[2]) return 1 " +
                    "end ";

            DefaultRedisScript<Long> redisScript = new DefaultRedisScript(script, Long.class);
            execute = redisTemplate.execute(redisScript, Arrays.asList(lua_key), String.valueOf(total), String.valueOf(seconds));
            if(execute == 0) {
                return "0";
            }

            // doSomething
            Thread.sleep(getSleep());

            return "1";
        } catch (Exception ex) {
            ex.printStackTrace();
        } finally {
            if (execute == 1) {
                redisTemplate.opsForValue().decrement(lua_key);
            }
        }
        return "0";
    }

    private int getRequestValue(String key) {
        String oldValueStr = redisTemplate.opsForValue().get(key);
        int oldValueInt = 0;
        if (!ObjectUtils.isEmpty(oldValueStr)) {
            try {
                oldValueInt = Integer.valueOf(oldValueInt);
            } catch (Exception ex) {
                ex.printStackTrace();
            }
        }
        return oldValueInt;
    }

    /**
     * 3.分布式锁模式
     */
    private String lock_key = "lock_key";
    @RequestMapping("lock")
    public String lock(){

        // 1.前置判断
        int oldValueInt = getRequestValue(lock_key);
        log.info("oldValueInt={}", oldValueInt);
        if (++oldValueInt > total) {
            return "0";
        }

        boolean requestFlag = false;
        try {
            // 2.分布式锁-内部修改限流参数
            String uuid = UUID.randomUUID().toString().replaceAll("-", "");
            try {
                boolean tryLock = tryLock(uuid);
                log.info("tryLock: key={},result={}", uuid, tryLock);
                if (!tryLock) {
                    return "0";
                }

                // 使用redis自增
                redisTemplate.opsForValue().increment(lock_key);
                // 不能使用原值自增
//                redisTemplate.opsForValue().set(lock_key, String.valueOf(oldValueInt));
                redisTemplate.expire(lock_key, seconds, TimeUnit.SECONDS);
                requestFlag = true;
            } catch (Exception ex) {
                ex.printStackTrace();
            } finally {
                boolean unlock = unlock(uuid);
                log.info("unlock: key={},result={}", uuid, unlock);
            }

            // doSomething
            Thread.sleep(getSleep());
            log.info("进入执行");
            return "1";
        } catch (Exception ex) {
            ex.printStackTrace();
        } finally {
            // 判断是否需要释放资源
            if (requestFlag) {
                // 递减
                redisTemplate.opsForValue().decrement(lock_key);
                // lua：判断+递减
//                String script =
//                        "local ret={} ret[1]=redis.call('get', KEYS[1]) " +
//                        "if tonumber(ret[1]) > 0 then " +
//                                "ret[2]=redis.call('decr', KEYS[1])  redis.call('expire', KEYS[1], ARGV[1])" +
//                        "end return ret";
//
//                DefaultRedisScript<List> redisScript = new DefaultRedisScript(script, List.class);
//                List<Long> execute = redisTemplate.execute(redisScript, Arrays.asList(lock_key), String.valueOf(seconds));
//                log.info("release={}", execute);
            }
        }
        return "0";
    }

    /**
     * 尝试设值，超时3s，则失败
     */
    private String redis_lock_key = "redis_lock_key";
    private boolean tryLock(String uuid) {
        try {
            long start = System.currentTimeMillis();

            while (true) {
                // set nx px
                Boolean setIfAbsent = redisTemplate.opsForValue().setIfAbsent(redis_lock_key, uuid, 3, TimeUnit.SECONDS);
                if (setIfAbsent) return true;

                long time = System.currentTimeMillis() - start;
                if (time >= 3000) {
                    return false;
                }
                Thread.sleep(1);
            }
        } catch (Exception ex) {
            ex.printStackTrace();
        }

        return true;
    }

    /**
     * 解锁使用lua
     * @return
     */
    private boolean unlock(String uuid) {
        String script =
            "if redis.call('get', KEYS[1]) == ARGV[1] then " +
                "return redis.call('del', KEYS[1]) " +
            "else return 0 " +
            "end ";

        DefaultRedisScript<Long> redisScript = new DefaultRedisScript(script, Long.class);
        Long execute = redisTemplate.execute(redisScript, Arrays.asList(redis_lock_key), uuid);
        return execute == 1;
    }

    /**
     * 4.redisson模式
     */
    private String redisson_lock_key = "redisson_lock_key";
    private String redisson_key = "redisson_key";
    @RequestMapping("redisson")
    public String redisson(){

        // 1.前置判断
        int oldValueInt = getRequestValue(lock_key);
        log.info("oldValueInt={}", oldValueInt);
        if (++oldValueInt > total) {
            return "0";
        }

        boolean requestFlag = false;
        try {
            // 2.RLock-内部修改限流参数
            RLock lock = redisson.getLock(redisson_lock_key);
            try {
                lock.lock();

                // 使用redis自增
                redisTemplate.opsForValue().increment(redisson_key);
                redisTemplate.expire(redisson_key, seconds, TimeUnit.SECONDS);
                requestFlag = true;
            } catch (Exception ex) {
                ex.printStackTrace();
            } finally {
                lock.unlock();
            }

            // doSomething
            Thread.sleep(getSleep());
            log.info("进入执行");
            return "1";
        } catch (Exception ex) {
            ex.printStackTrace();
        } finally {
            if (requestFlag) {
                redisTemplate.opsForValue().decrement(redisson_key);
            }
        }
        return "0";
    }

}
